package com.lightboat.common.algorithm;

import lombok.Getter;
import lombok.Setter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;

/**
 * 基于前缀和的权重轮盘抽取算法
 */
public class WeightedLotteryWheel<T> {

    private final List<Prize<T>> prizes;
    private int[] prefixSum;  // 前缀和数组
    private int totalWeight;
    public WeightedLotteryWheel(List<Prize<T>> prizes) {
        this.prizes = new ArrayList<>(prizes);
        this.buildPrefixSum();
    }

    /**
     * 构建前缀和数组 - 预计算阶段
     * 时间复杂度：O(n)
     * 空间复杂度：O(n)
     */
    private void buildPrefixSum() {
        int n = prizes.size();
        prefixSum = new int[n];
        prefixSum[0] = prizes.get(0).weight;

        for (int i = 1; i < n; i++) {
            prefixSum[i] = prefixSum[i - 1] + prizes.get(i).weight;
        }

        totalWeight = prefixSum[n - 1];
    }

    /**
     * 执行抽取 - 运行时阶段
     * 时间复杂度：O(log n) - 使用二分查找
     * 空间复杂度：O(1)
     */
    public Prize<T> drawPrize() {
        if (totalWeight == 0) {
            return null;
        }

        // 生成 [1, totalWeight] 范围内的随机数
        int randomValue = ThreadLocalRandom.current().nextInt(1, totalWeight + 1);

        // 二分查找定位抽取数据
        int index = binarySearch(randomValue);
        return prizes.get(index);
    }

    /**
     * 二分查找第一个大于等于target的位置
     */
    private int binarySearch(int target) {
        int left = 0, right = prefixSum.length - 1;
        int result = 0;

        while (left <= right) {
            int mid = left + (right - left) / 2;
            if (prefixSum[mid] >= target) {
                result = mid;
                right = mid - 1;
            } else {
                left = mid + 1;
            }
        }

        return result;
    }

    // 抽取数据信息
    @Getter
    @Setter
    public static class Prize<T> {
        private String id;
        private String name;
        private int weight;  // 权重
        private T data;

        public Prize(String id, String name, int weight) {
            this.id = id;
            this.name = name;
            this.weight = weight;
        }

        public Prize(String id, String name, int weight, T data) {
            this.id = id;
            this.name = name;
            this.weight = weight;
            this.data = data;
        }
    }
}